from typing import List, Dict, Optional

# Notice 前提：配置好系统变量
from backend.app.database.db_pinecone import rdb_conn


class DbPinecone(object):
    """矢量数据库"""
    _namespace = ""
    _index_name = ""

    def __init__(self, namespace: str = "", index_name: str = "avl"):
        self._namespace = namespace
        self._index_name = index_name

    def get_rds(self):
        """获取索引连接"""
        return rdb_conn.get_rds(index_name=self._index_name, namespace=self._namespace)

    @staticmethod
    def get_embeddings():
        return rdb_conn.get_embeddings()

    def _get_rds_index(self):
        index_name = self._index_name
        rds = self.get_rds()
        return rds.get_pinecone_index(index_name)

    @staticmethod
    def format_schema_data(result_list):
        meta_list = []
        for item in result_list:
            x = item["metadata"]
            x["vector_id"] = item["id"]
            meta_list.append(x)
        return meta_list

    def get_by_ids(self, vector_ids: List[str]):
        """根据向量ID获取列表"""
        if not vector_ids:
            return
        rds_index = self._get_rds_index()
        result = rds_index.fetch(vector_ids)
        if not (result and result.get("vectors")):
            return

        return self.format_schema_data(result["vectors"].values())

    def get_by_goods_id(self, device_name: str, goods_id: int):
        rds_index = self._get_rds_index()
        embeddings = self.get_embeddings()
        vector = embeddings.embed_query(device_name)
        conditions = dict()
        if goods_id and goods_id > 0:
            conditions["g_id"] = {"$eq": goods_id}

        res = rds_index.query(
            namespace=self._namespace,
            index_name=self._index_name,
            vector=vector,
            filter=conditions,
            include_values=True,
            top_k=1,
            include_metadata=True,
        )
        return res["matches"][0] if res else None

    def get_by_text(self, text: str, g_id: int, include_values: bool = False,
                    include_metadata: bool = True, **kwargs):
        rds_index = self._get_rds_index()
        embeddings = self.get_embeddings()
        vector = embeddings.embed_query(text)
        conditions = dict()
        if g_id and g_id > 0:
            conditions["g_id"] = {"$eq": g_id}

        res = rds_index.query(
            namespace=self._namespace,
            index_name=self._index_name,
            vector=vector,
            filter=conditions,
            include_values=include_values,
            top_k=3,
            include_metadata=include_metadata,
            **kwargs
        )
        if not (res or res.get("matches")):
            return
        return self.format_schema_data(res["matches"])

    def update_by_id(self, vector_id: str, text: str, metadata: Dict, **kwargs):
        vector_list = self.get_embeddings().embed_query(text)
        rds_index = self._get_rds_index()
        return rds_index.update(
            namespace=self._namespace,
            index_name=self._index_name,
            id=vector_id,
            values=vector_list,
            set_metadata=metadata,
            **kwargs
        )

    def batch_add(self, texts: List[str], metadata: List[Dict], **kwargs) -> List[str]:
        """批量插入数据"""
        index_name = self._index_name
        namespace = self._namespace
        rdb = self.get_rds()
        return rdb.add_texts(texts, metadatas=metadata, index_name=index_name, namespace=namespace, **kwargs)

    def delete_by_ids(self, vector_ids: List[str], gids: List[int] = None) -> None:
        """检测是否有效"""
        if not vector_ids:
            return
        rdb = self.get_rds()
        conditions = dict()
        if gids:
            conditions["g_id"] = {"$in": gids}
        return rdb.delete(ids=vector_ids, filter=conditions, namespace=self._namespace, index_name=self._index_name)

    async def abatch_add(self, texts: List[str], metadata: List[Dict], **kwargs) -> List[str]:
        """批量插入数据-异步"""
        # 判断是否为空
        if not (texts and metadata):
            return []
        index_name = self._index_name
        namespace = self._namespace
        rdb = self.get_rds()
        return await rdb.aadd_texts(texts, metadatas=metadata, index_name=index_name, namespace=namespace, **kwargs)

    def batch_upsert(self, data: List[Dict]):
        """
        批量更新
        :param data: [{"metadata": 设备属性字典, "id": 向量ID, "values": 向量}]
        :return: {'upserted_count': 修改条数}
        """
        if not data:
            return
        embd = self.get_embeddings()
        rds_index = self._get_rds_index()
        lists = []
        for md in data:
            values = md["values"]
            if type(values) == str:
                values = embd.embed_query(md["values"])
                md["metadata"]["text"] = md["values"]

            item = dict(metadata=md["metadata"], id=md["id"], values=values)
            lists.append(item)
        return rds_index.upsert(lists)

    async def adelete_by_ids(self, vector_ids: List[str], gids: List[int] = None) -> Optional[bool]:
        """删除-异步"""
        if not vector_ids:
            return
        rdb = self.get_rds()
        conditions = dict()
        if gids:
            conditions["g_id"] = {"$in": gids}
        return await rdb.adelete(ids=vector_ids, filter=conditions, namespace=self._namespace,
                                 index_name=self._index_name)

